"""基础矩阵运算模块
包含矩阵加减乘除、转置、逆矩阵、幂运算、迹、行列式等基础运算
支持复数矩阵、稀疏矩阵等多种数据类型
"""
import numpy as np
from scipy import sparse
from typing import Union, Optional, Tuple
import warnings
from debug_tools import debug_matrix_operation

MatrixType = Union[np.ndarray, sparse.spmatrix]

class BasicMatrixOperations:
    """基础矩阵运算类"""
    
    @staticmethod
    @debug_matrix_operation("matrix_add")
    def matrix_add(a: MatrixType, b: MatrixType) -> MatrixType:
        """
        矩阵加法
        :param a: 矩阵A
        :param b: 矩阵B
        :return: A + B
        """
        # 输入验证
        if a is None or b is None:
            raise ValueError("输入矩阵不能为None")
        
        # 维度检查
        if hasattr(a, 'shape') and hasattr(b, 'shape'):
            if a.shape != b.shape:
                raise ValueError(f"矩阵维度必须匹配: {a.shape} vs {b.shape}")
        
        # 数值检查
        if not sparse.issparse(a) and not np.isfinite(a).all():
            raise ValueError("矩阵A包含无效数值(NaN或Inf)")
        if not sparse.issparse(b) and not np.isfinite(b).all():
            raise ValueError("矩阵B包含无效数值(NaN或Inf)")
        
        try:
            if sparse.issparse(a) or sparse.issparse(b):
                result = a + b
            else:
                result = np.add(a, b)
            
            # 结果验证
            if not sparse.issparse(result) and not np.isfinite(result).all():
                raise ValueError("矩阵加法结果包含无效数值")
            
            return result
        except Exception as e:
            raise ValueError(f"矩阵加法计算失败: {str(e)}")
    
    @staticmethod
    def matrix_subtract(a: MatrixType, b: MatrixType) -> MatrixType:
        """
        矩阵减法
        :param a: 矩阵A
        :param b: 矩阵B
        :return: A - B
        """
        # 验证矩阵维度
        if hasattr(a, 'shape') and hasattr(b, 'shape'):
            if a.shape != b.shape:
                raise ValueError(f"Matrix dimensions must match for subtraction: {a.shape} vs {b.shape}")
        
        if sparse.issparse(a) or sparse.issparse(b):
            return a - b
        return np.subtract(a, b)
    
    @staticmethod
    def matrix_multiply(a: MatrixType, b: MatrixType) -> MatrixType:
        """
        矩阵乘法
        :param a: 矩阵A
        :param b: 矩阵B
        :return: A @ B
        """
        # 输入验证
        if a is None or b is None:
            raise ValueError("输入矩阵不能为None")
        
        # 维度检查
        if hasattr(a, 'shape') and hasattr(b, 'shape'):
            if len(a.shape) < 2 or len(b.shape) < 2:
                raise ValueError(f"矩阵必须至少是二维的: A{a.shape}, B{b.shape}")
            if a.shape[-1] != b.shape[-2]:
                raise ValueError(f"矩阵乘法维度不匹配: A{a.shape} @ B{b.shape}")
        
        # 数值检查
        if not sparse.issparse(a) and not np.isfinite(a).all():
            raise ValueError("矩阵A包含无效数值(NaN或Inf)")
        if not sparse.issparse(b) and not np.isfinite(b).all():
            raise ValueError("矩阵B包含无效数值(NaN或Inf)")
        
        # 大小检查（防止内存溢出）
        if hasattr(a, 'shape') and hasattr(b, 'shape'):
            result_size = a.shape[0] * b.shape[1]
            if result_size > 1e8:  # 超过1亿元素
                raise ValueError(f"结果矩阵过大，可能导致内存不足: {a.shape[0]}×{b.shape[1]}")
        
        try:
            if sparse.issparse(a) or sparse.issparse(b):
                result = a @ b
            else:
                result = np.matmul(a, b)
            
            # 结果验证
            if not sparse.issparse(result) and not np.isfinite(result).all():
                raise ValueError("矩阵乘法结果包含无效数值")
            
            return result
        except MemoryError:
            raise MemoryError("矩阵乘法导致内存不足")
        except Exception as e:
            raise ValueError(f"矩阵乘法计算失败: {str(e)}")
    
    @staticmethod
    def matrix_elementwise_multiply(a: MatrixType, b: MatrixType) -> MatrixType:
        """
        矩阵逐元素乘法（Hadamard积）
        :param a: 矩阵A
        :param b: 矩阵B
        :return: A ⊙ B
        """
        if sparse.issparse(a) or sparse.issparse(b):
            return a.multiply(b)
        return np.multiply(a, b)
    
    @staticmethod
    def matrix_divide(a: MatrixType, b: Union[MatrixType, float, int]) -> MatrixType:
        """
        矩阵除法（右乘逆矩阵或标量除法）
        :param a: 矩阵A
        :param b: 矩阵B或标量
        :return: A / B 或 A * B^(-1)
        """
        if isinstance(b, (int, float, complex)):
            if sparse.issparse(a):
                return a / b
            return np.divide(a, b)
        else:
            # 矩阵除法：A / B = A * B^(-1)
            return BasicMatrixOperations.matrix_multiply(a, BasicMatrixOperations.matrix_inverse(b))
    
    @staticmethod
    def matrix_transpose(a: MatrixType) -> MatrixType:
        """
        矩阵转置
        :param a: 输入矩阵
        :return: A^T
        """
        if sparse.issparse(a):
            return a.T
        return np.transpose(a)
    
    @staticmethod
    def matrix_conjugate_transpose(a: MatrixType) -> MatrixType:
        """
        矩阵共轭转置（厄米转置）
        :param a: 输入矩阵
        :return: A^H
        """
        if sparse.issparse(a):
            return a.H
        return np.conj(a).T
    
    @staticmethod
    def matrix_inverse(a: np.ndarray, method: str = 'lu') -> np.ndarray:
        """
        矩阵求逆
        :param a: 输入方阵
        :param method: 求逆方法 ('lu', 'cholesky', 'svd')
        :return: A^(-1)
        """
        # 输入验证
        if a is None:
            raise ValueError("输入矩阵不能为None")
        
        if not isinstance(a, np.ndarray):
            raise ValueError("输入必须是numpy数组")
        
        # 方阵检查
        if a.shape[0] != a.shape[1]:
            raise ValueError(f"矩阵必须是方阵才能求逆: {a.shape}")
        
        # 数值检查
        if not np.isfinite(a).all():
            raise ValueError("矩阵包含无效数值(NaN或Inf)")
        
        # 大小检查
        if a.shape[0] > 1000:
            raise ValueError(f"矩阵过大，可能导致计算不稳定: {a.shape}")
        
        # 奇异性检查
        det_a = np.linalg.det(a)
        if abs(det_a) < 1e-12:
            raise ValueError(f"矩阵接近奇异，无法求逆，行列式值: {det_a:.2e}")
        
        # 条件数检查
        cond_a = np.linalg.cond(a)
        if cond_a > 1e12:
            raise ValueError(f"矩阵条件数过大，可能导致数值不稳定: {cond_a:.2e}")
        
        # 方法验证
        valid_methods = ['lu', 'cholesky', 'svd']
        if method not in valid_methods:
            raise ValueError(f"无效的求逆方法: {method}，支持的方法: {valid_methods}")
        
        try:
            if method == 'lu':
                result = np.linalg.inv(a)
            elif method == 'cholesky':
                # 检查对称性
                if not np.allclose(a, a.T, atol=1e-10):
                    warnings.warn("矩阵不对称，Cholesky分解可能失败，回退到LU分解")
                    result = np.linalg.inv(a)
                else:
                    # 检查正定性
                    eigenvals = np.linalg.eigvals(a)
                    if not np.all(eigenvals > 1e-12):
                        warnings.warn("矩阵不是正定的，Cholesky分解失败，回退到LU分解")
                        result = np.linalg.inv(a)
                    else:
                        try:
                            L = np.linalg.cholesky(a)
                            result = np.linalg.inv(L.T) @ np.linalg.inv(L)
                        except np.linalg.LinAlgError:
                            warnings.warn("Cholesky分解失败，回退到LU分解")
                            result = np.linalg.inv(a)
            elif method == 'svd':
                # 使用SVD求伪逆
                result = np.linalg.pinv(a)
            
            # 结果验证
            if not np.isfinite(result).all():
                raise ValueError("矩阵求逆结果包含无效数值")
            
            # 逆矩阵正确性验证
            identity_check = a @ result
            identity_error = np.linalg.norm(identity_check - np.eye(a.shape[0]))
            if identity_error > 1e-10:
                raise ValueError(f"矩阵求逆不准确，验证误差: {identity_error:.2e}")
            
            return result
            
        except np.linalg.LinAlgError as e:
            raise ValueError(f"矩阵求逆线性代数错误: {str(e)}")
        except Exception as e:
            raise ValueError(f"矩阵求逆计算失败: {str(e)}")
    
    @staticmethod
    def matrix_pseudoinverse(a: np.ndarray, rcond: Optional[float] = None) -> np.ndarray:
        """
        矩阵伪逆（Moore-Penrose逆）
        :param a: 输入矩阵
        :param rcond: 奇异值截断阈值
        :return: A^+
        """
        return np.linalg.pinv(a, rcond=rcond)
    
    @staticmethod
    def matrix_power(a: np.ndarray, n: Union[int, float]) -> np.ndarray:
        """
        矩阵幂运算
        :param a: 输入方阵
        :param n: 幂次（整数或实数）
        :return: A^n
        """
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square for power operation")
        
        if isinstance(n, int):
            if n == 0:
                return np.eye(a.shape[0], dtype=a.dtype)
            elif n > 0:
                result = np.eye(a.shape[0], dtype=a.dtype)
                base = a.copy()
                while n > 0:
                    if n % 2 == 1:
                        result = result @ base
                    base = base @ base
                    n //= 2
                return result
            else:  # n < 0
                return BasicMatrixOperations.matrix_power(BasicMatrixOperations.matrix_inverse(a), -n)
        else:
            # 实数幂：使用特征值分解
            eigenvals, eigenvecs = np.linalg.eig(a)
            if np.any(eigenvals.real < 0) and not np.allclose(eigenvals.imag, 0):
                warnings.warn("Matrix has negative eigenvalues, result may be complex")
            
            powered_eigenvals = np.power(eigenvals, n)
            return eigenvecs @ np.diag(powered_eigenvals) @ np.linalg.inv(eigenvecs)
    
    @staticmethod
    def matrix_trace(a: MatrixType) -> Union[float, complex]:
        """
        矩阵迹（对角元素之和）
        :param a: 输入方阵
        :return: tr(A)
        """
        if sparse.issparse(a):
            return a.diagonal().sum()
        return np.trace(a)
    
    @staticmethod
    def matrix_determinant(a: np.ndarray) -> Union[float, complex]:
        """
        矩阵行列式
        :param a: 输入方阵
        :return: det(A)
        """
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square for determinant calculation")
        return np.linalg.det(a)
    
    @staticmethod
    def matrix_rank(a: np.ndarray, tol: Optional[float] = None) -> int:
        """
        矩阵秩
        :param a: 输入矩阵
        :param tol: 容差
        :return: rank(A)
        """
        return np.linalg.matrix_rank(a, tol=tol)
    
    @staticmethod
    def matrix_norm(a: MatrixType, ord: Union[str, int, float] = 'fro') -> float:
        """
        矩阵范数
        :param a: 输入矩阵
        :param ord: 范数类型 ('fro', 1, 2, -1, -2, np.inf, -np.inf, 'nuc')
        :return: ||A||_ord
        """
        if sparse.issparse(a):
            if ord == 'fro':
                return sparse.linalg.norm(a, ord='fro')
            else:
                return sparse.linalg.norm(a, ord=ord)
        return np.linalg.norm(a, ord=ord)
    
    @staticmethod
    def matrix_condition_number(a: np.ndarray, p: Union[str, int, float] = 2) -> float:
        """
        矩阵条件数
        :param a: 输入矩阵
        :param p: 范数类型
        :return: cond(A)
        """
        return np.linalg.cond(a, p=p)
    
    @staticmethod
    def create_identity_matrix(n: int, dtype: type = float) -> np.ndarray:
        """
        创建单位矩阵
        :param n: 矩阵维度
        :param dtype: 数据类型
        :return: n×n单位矩阵
        """
        return np.eye(n, dtype=dtype)
    
    @staticmethod
    def create_zeros_matrix(m: int, n: int, dtype: type = float) -> np.ndarray:
        """
        创建零矩阵
        :param m: 行数
        :param n: 列数
        :param dtype: 数据类型
        :return: m×n零矩阵
        """
        return np.zeros((m, n), dtype=dtype)
    
    @staticmethod
    def create_ones_matrix(m: int, n: int, dtype: type = float) -> np.ndarray:
        """
        创建全1矩阵
        :param m: 行数
        :param n: 列数
        :param dtype: 数据类型
        :return: m×n全1矩阵
        """
        return np.ones((m, n), dtype=dtype)
    
    @staticmethod
    def create_diagonal_matrix(diag_elements: np.ndarray) -> np.ndarray:
        """
        创建对角矩阵
        :param diag_elements: 对角元素
        :return: 对角矩阵
        """
        return np.diag(diag_elements)
    
    @staticmethod
    def matrix_kronecker_product(a: np.ndarray, b: np.ndarray) -> np.ndarray:
        """
        Kronecker积
        :param a: 矩阵A
        :param b: 矩阵B
        :return: A ⊗ B
        """
        return np.kron(a, b)
    
    @staticmethod
    def matrix_frobenius_inner_product(a: np.ndarray, b: np.ndarray) -> Union[float, complex]:
        """
        Frobenius内积
        :param a: 矩阵A
        :param b: 矩阵B
        :return: <A, B>_F = tr(A^H B)
        """
        return np.trace(np.conj(a).T @ b)
    
    @staticmethod
    def matrix_commutator(a: np.ndarray, b: np.ndarray) -> np.ndarray:
        """
        矩阵交换子
        :param a: 矩阵A
        :param b: 矩阵B
        :return: [A, B] = AB - BA
        """
        return a @ b - b @ a
    
    @staticmethod
    def matrix_anticommutator(a: np.ndarray, b: np.ndarray) -> np.ndarray:
        """
        矩阵反交换子
        :param a: 矩阵A
        :param b: 矩阵B
        :return: {A, B} = AB + BA
        """
        return a @ b + b @ a